package com.iot.ruler.core;

import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 用于缓存线程池无法处理的任务
 * 当线程池执行拒绝策略时，会将无法执行的任务缓存到队列中
 * 实时监控线程池的运行状态，当线程池队列有空闲时取出缓存任务执行
 *
 * @author lzb
 * @Date 2023/2/10 17:04
 */
@Slf4j
public class CacheRejectHandler implements RejectedExecutionHandler, OnThreadExecuteListener {

    /**
     * 缓存任务队列
     */
    private final LinkedBlockingQueue<Runnable> cacheQueue = new LinkedBlockingQueue<>();
    /**
     * 控制缓存任务出队执行的线程池
     */
    private final ThreadPoolExecutor cacheExecutor = new ThreadPoolExecutor(
            2, 2,
            0, TimeUnit.SECONDS,
            new ArrayBlockingQueue<>(1),
            new ThreadFactory() {
                private final AtomicInteger atomicInteger = new AtomicInteger(0);

                @Override
                public Thread newThread(Runnable r) {
                    Thread t = new Thread(r);
                    t.setPriority(Thread.MAX_PRIORITY);
                    t.setDaemon(true);
                    t.setName("CacheRejectHandler:" + atomicInteger.incrementAndGet());
                    System.out.println("创建拒绝策略处理线程：" + t.getName());
                    return t;
                }
            });
    private final LinkedBlockingQueue<Integer> lockOutQueue = new LinkedBlockingQueue<>();
    private final LinkedBlockingQueue<Integer> lockSwitchQueue = new LinkedBlockingQueue<>();
    private boolean isTerminated = false;
    /**
     * 用来执行缓存任务的线程池
     */
    private ThreadPoolExecutor taskExecutor;
    /**
     * 缓存任务出队送入线程池继续执行
     */
    private final Runnable cacheOutTask = new Runnable() {
        @Override
        public void run() {
            log.info("Ruler任务缓存启动");
            do {
                // 如果lockQueue里无数据则会阻塞线程执行
                try {
                    lockOutQueue.take();
                    Runnable newTask = cacheQueue.take();
                    if (taskExecutor.getQueue().remainingCapacity() > 0) {
                        // 取出缓冲的队列数据，重新放入线程池执行
                        taskExecutor.execute(newTask);
                        continue;
                    }
                    cacheQueue.put(newTask);
                } catch (Exception e) {
                    log.error("cacheOutTask", e);
                }
            } while (!isTerminated);
        }
    };
    /**
     * 缓存任务出队送入线程池继续执行
     */
    private final Runnable cacheSwitchTask = new Runnable() {
        @Override
        public void run() {
            log.info("Ruler任务缓存重放控制启动");
            do {
                try {
                    // 如果lockSwitchQueue里无数据则会阻塞线程执行
                    lockSwitchQueue.take();
                    int remainingCapacity = taskExecutor.getQueue().remainingCapacity();
                    do {
                        lockOutQueue.put(1);
                        remainingCapacity--;
                    } while (remainingCapacity > 0);
                } catch (Exception e) {
                    log.error("cacheSwitchTask", e);
                }
            } while (!isTerminated);
        }
    };

    public CacheRejectHandler() {

    }

    /**
     * 绑定线程池并开启监控
     */
    public void bind(ThreadPoolExecutor executor) {
        isTerminated = false;
        this.taskExecutor = executor;
        cacheExecutor.execute(cacheOutTask);
        cacheExecutor.execute(cacheSwitchTask);
    }

    public void offer(Runnable runnable) {
        cacheQueue.offer(runnable);
    }

    @Override
    public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
        offer(r);
    }

    @Override
    public void beforeExecute(Thread t, Runnable r) {

    }

    @SneakyThrows
    @Override
    public void afterExecute(Runnable r, Throwable t) {
        if (cacheQueue.isEmpty()) {
            // 没有缓冲任务则忽略
            return;
        }
        // 还有缓冲任务队列不为空(有任务没执行)
        // 线程池的任务队列已被消耗空，即恢复为0时
        // 将缓冲任务重放任务锁打开，重新放入线程池执行
        // 一次最多释放线程池队列的剩余容量(初始容量)
        if (taskExecutor.getQueue().isEmpty()) {
            System.out.println("恢复switch任务");
            lockSwitchQueue.put(1);
        }
    }

    @Override
    public void terminated() {
        cacheExecutor.shutdown();
        isTerminated = true;
    }
}
